import asyncio
import base64
import json
import os.path
from pathlib import Path

import cv2
import nanoid
import numpy as np
import tornado

from tasks.Manager import TaskManager

training_tasks = {}
tensorboard_tasks = {}
tm = TaskManager()




class Server(tornado.web.RequestHandler):

    def initialize(self, config):
        self.config = config



    async def post(self):
        cmd_body = json.loads(self.request.body)
        command = cmd_body["cmd"]
        params = cmd_body["params"]
        resp = {}
        resp["status"] = 200
        resp["msg"] = ""

        if command == "viewLabel":
            src_path = "E:/data/CrowdHuman20250110/images/train/273278,32b49000523ef483.jpg"
            label_path = Path("E:/data/CrowdHuman20250110/annotation_train.odgt")
            filename = Path(src_path).stem
            src_img_cv = cv2.imread(src_path)
            with open(label_path,"r") as file:
                for line in file:
                    line = line.strip()
                    labels = json.loads(line)
                    if labels["ID"] == filename:
                        boxes = labels["gtboxes"]
                        for box in boxes:
                            if box["tag"] == "mask" or box.get("extra",{}).get("ignore",0)==1:
                                continue
                            name = box["tag"]
                            points = box["fbox"]
                            cv2.rectangle(src_img_cv,(points[0],points[1]),(points[0]+points[2],points[1]+points[3]),[220, 20, 60],2)
                            cv2.imwrite(f"{filename}-output.png",src_img_cv)
            json_str = json.dumps(resp)
            self.set_header("Content-Type", "application/json")
            self.write(json_str)

        if command == "train":
            params["train_log_path"] = self.config["app"]["train"]["logPath"]
            params["pretrain_path"] = self.config["app"]["train"]["pretrainPath"]
            task_id = tm.create_training_task(params)
            self.set_header("Content-Type", "application/json")
            resp["msg"] = f"training task {task_id} is running"
            json_str = json.dumps(resp)
            self.write(json_str)


        if command == "cv_predict_image":
            params["colors"] = self.config["app"]["colors"]
            result_img,result = tm.cv_predict_image(params)
            ret_type = params["ret_type"]
            if ret_type == "image":
                self.set_header("Content-Type", "image/jpeg")
                with open(result_img,"rb") as f:
                    self.write(f.read())
                    f.close()
            if ret_type == "not_image":
                self.set_header("Content-Type", "application/json")
                resp["result"] = result
                json_str = json.dumps(resp)
                self.write(json_str)
            os.remove(result_img)


        if command == "deploy":
            params = cmd_body["params"]
            status,msg = tm.deploy_model(params)
            resp["status"] = status
            resp["msg"] = msg
            json_str = json.dumps(resp)
            self.write(json_str)


        if command == "base64encode":
            body = {}
            params = {}
            img_path = cmd_body["params"]["img_path"]
            with open(img_path,"rb") as f:
                img_data = f.read()
            img_base64 = base64.b64encode(img_data).decode("utf-8")
            body["cmd"]="cv_predict_image"
            params["ret_type"] = "image"
            params["task_type"] = "object_detection"
            params["image"]=img_base64
            params["model_id"] = "yolo11n-WsVZ_taCesoohBv4daeJq"
            body["params"] = params
            with open("predictBody.json","w") as f2:
                f2.write(json.dumps(body))



